import numpy as np
import cv2 as cv
import glob
import os
import platform


def calc_disp_bm(main_, aux_):
    max_disp = 32
    win_size = 9
    uniqueness = 30
    im_L = main_
    im_R = aux_
    stereo = cv.StereoBM_create(numDisparities=max_disp, blockSize=win_size)
    stereo.setUniquenessRatio(uniqueness)
    disp_s16 = stereo.compute(im_L, im_R)
    disp_s16 = np.maximum(disp_s16, 0)
    disp_f32 = np.array(disp_s16, dtype=np.float32)
    disp_u8 = np.array(disp_f32 * 255.0 / (16.0*max_disp), dtype=np.uint8)
    #disp_u8 = np.minimum(disp_u8, int(real_max_disp*255 / max_disp))
    return disp_u8
    

def calc_disp_sgbm(main_, aux_):
    max_disp = 32
    win_size = 9
    uniqueness = 30
    im_L = main_
    im_R = aux_
    stereo = cv.StereoSGBM_create(
            minDisparity=0,
            numDisparities=max_disp,
            blockSize=win_size,
            P1=24*win_size*win_size,
            P2=96*win_size*win_size,
            uniquenessRatio=uniqueness,
            preFilterCap=63,
            mode=3)
    disp_s16 = stereo.compute(im_L, im_R)
    disp_s16 = np.maximum(disp_s16, 0)
    disp_f32 = np.array(disp_s16, dtype=np.float32)
    disp_u8 = np.array(disp_f32 * 255.0 / (16.0*max_disp), dtype=np.uint8)
    #disp_u8 = np.minimum(disp_u8, int(real_max_disp*255 / max_disp))
    return disp_u8


def main_procedure(in_path, out_path, training: bool):
    begin_idx = 0
    if os.path.exists(out_path):
        files_rgbd = glob.glob(out_path + '/*.PNG')
        for i in range(len(files_rgbd)):
            id_ = int(files_rgbd[i][-12:-4]) + 1
            begin_idx = max(id_, begin_idx)
    else:
        os.mkdir(out_path)
    print(begin_idx)

    if training:
        files_main = glob.glob(in_path + '/*_L.PNG')
        files_aux = glob.glob(in_path + '/*_R.PNG')
    else:
        files_main = glob.glob(in_path + '/*main.PNG')
        files_aux = glob.glob(in_path + '/*aux.PNG')
    assert len(files_main) == len(files_aux)
    crop = 32
    for i in range(len(files_main)):
        print(files_main[i])
        im_main = cv.imread(files_main[i], 1)
        im_aux = cv.imread(files_aux[i], 1)
        main_grey = cv.cvtColor(im_main, cv.COLOR_BGR2GRAY)
        aux_grey = cv.cvtColor(im_aux, cv.COLOR_BGR2GRAY)
        # real_max_disp = int(files_main[i][-10:-8])
        disp = calc_disp_sgbm(main_grey, aux_grey)
        # RGB + D ->RGBD
        rgbd = np.stack([im_main[:, crop:, 0],
                         im_main[:, crop:, 1],
                         im_main[:, crop:, 2],
                         disp[:, crop:]], axis=-1)
        cv.imwrite(out_path + '/%08d.PNG' % (i + begin_idx), rgbd)
        # cv.imwrite(out_path + '/%08d.PNG'%(i+begin_idx), disp[:, crop:])


def main():
    in_path = 'F:/dump/Matting/PORTRAIT'
    out_path = 'F:/dump/Matting/RGBD'
    main_procedure(in_path, out_path, training=True)


if __name__ == '__main__':
    main()

